"""
visualization
"""
import os
import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def r2_score(label, predict):
    """
    计算R2分数（决定系数），衡量模型预测效果
    R2=1表示完美预测，R2=0表示模型预测效果与均值相当
    
    Args:
        label: 真实标签值
        predict: 模型预测值
        
    Returns:
        float: R2分数
    """
    label_mean = np.mean(label)
    ss_tot = np.sum((label - label_mean) ** 2)  # 总离差平方和
    ss_res = np.sum((label - predict) ** 2)    # 残差平方和
    r2_value = 1.0 - (ss_res / ss_tot)         # R2分数计算公式
    return r2_value


def mean_squared_error(label, predict):
    """
    计算均方误差(MSE)，衡量预测值与真实值的偏差
    MSE越小表示预测效果越好
    
    Args:
        label: 真实标签值
        predict: 模型预测值
        
    Returns:
        float: 均方误差
    """
    return np.mean((label - predict) ** 2, axis=0)


def plt_loss_func(epochs, data, filename, is_train=True, prefix="result_picture"):
    """
    绘制损失函数变化曲线，标记最大值和最小值点
    
    Args:
        epochs: 训练轮次总数
        data: 损失值列表
        filename: 图片文件名
        is_train: 是否为训练损失（True）或测试损失（False）
        prefix: 图片保存路径前缀
    """
    plt.cla()  # 清除当前坐标轴
    max_value = max(data)
    min_value = min(data)
    max_index = data.index(max_value)
    min_index = data.index(min_value)

    plt.plot(list(range(epochs)), data)
    plt.plot(max_index, max_value, 'ks')  # 标记最大值点（黑色正方形）
    plt.plot(min_index, min_value, 'ks')  # 标记最小值点（黑色正方形）
    
    # 标注最大值点坐标
    show_max = '(' + str(max_index) + ',' + str(max_value) + ')'
    plt.annotate(show_max, xytext=(max_index, max_value), xy=(max_index, max_value))
    
    # 标注最小值点坐标（偏移10个单位避免遮挡）
    show_min = f"({min_index-10}, {min_value:0.1e})"
    plt.annotate(show_min, xytext=(min_index-10, min_value), xy=(min_index-10, min_value))
    
    plt.xlabel("epochs")
    plt.ylabel("train_loss" if is_train else "test_loss")
    plt.title("loss")

    file_name = os.path.join(prefix, filename)
    plt.savefig(file_name)


def plt_error_dis_analyze(config, data, filepath):
    """
    分析误差与壁面距离的关系，绘制两类直方图：
    1. 不同壁面距离区间的平均误差
    2. 不同误差区间的平均壁面距离
    
    Args:
        config: 配置字典（包含绘图参数）
        data: 包含误差和壁面距离的数据框
        filepath: 图片保存路径
    """
    plt.cla()
    label = data["post_label"].values
    predict = data["post_predict"].values
    error = np.abs(label - predict) / label  # 计算相对误差
    dis = data["dis"].values
    new_data = pd.DataFrame({"error": error, "dis": dis}).sort_values(by="error")
    new_data.to_csv(filepath + '.csv')  # 保存分析数据
    plt.figure(figsize=(8, 10), dpi=200)

    # 上半部分：0.01 < 壁面距离 < 0.2 时的误差-距离平均直方图
    plt.subplot(2, 1, 1)
    hist, bin_value = get_mean_error_by_condition(config["figure_bins"],
                                                  config["plt_dis_internal"], new_data)
    plt.xticks(range(0, 9), format_sci(bin_value[1:]), rotation=30)
    plt.bar(range(0, 9), hist[1:])
    for value, prob in zip(range(0, 9), hist[1:]):
        plt.text(value, prob + 0.05, f'{prob:0.0f}', size=8, ha='center', va='bottom')
    plt.xlabel('Dis', fontsize=10)
    plt.ylabel('Error', fontsize=10)
    plt.title('Error-Dis Average Hist  (0.02 < Dis < 0.2)')

    # 下半部分：0 < 误差 < 1 时的距离-误差平均直方图
    plt.subplot(2, 1, 2)
    hist, bin_value = get_mean_dis_by_condition(config["figure_bins"],
                                                config["plt_error_internal"], new_data)
    plt.xticks(range(0, 10), format_sci(bin_value), rotation=30)
    plt.bar(range(0, 10), hist)
    for value, prob in zip(range(0, 10), hist):
        plt.text(value, prob, f'{prob:0.2f}', size=8, ha='center', va='bottom')
    plt.xlabel('Error', fontsize=10)
    plt.ylabel('Dis', fontsize=10)
    plt.title('Dis-Error Average Hist (0 < Error < 1) ')

    plt.tight_layout()
    plt.savefig(filepath)
    plt.clf()


def plt_error_distribute(config, x_value, post_predict, post_label, dis, model_predict,
                         model_label, prefix, train_or_test="train"):
    """
    绘制不同区域的误差分布：
    1. 全部数据
    2. 近壁面区域
    3. 激波区域
    4. 误差与距离的关系分析图
    
    Args:
        config: 配置字典
        x_value: x坐标值
        post_predict: 后处理预测值
        post_label: 后处理真实值
        dis: 壁面距离
        model_predict: 模型输出预测值
        model_label: 模型输出真实值
        prefix: 保存路径前缀
        train_or_test: "train"或"test"标记
    """
    near_wall = train_or_test + '_near_wall'
    all_data = train_or_test + '_all'
    shock_wave = train_or_test + 'shock_wave'
    plt.cla()
    
    # 创建包含所有数据的数据框
    data = pd.DataFrame({"post_label": post_label, "post_predict": post_predict,
                         "model_label": model_label, "model_predict": model_predict,
                         "dis": dis, "X": x_value})

    filepath = os.path.join(prefix, all_data + '_error.png')
    plt_hist_and_line(data, filepath)  # 绘制全部数据误差图

    # 近壁面区域（dis < 配置值）
    data_near_wall = data[data["dis"] < config["near_wall_dis"]]
    filepath = os.path.join(prefix, near_wall + '_error.png')
    plt_hist_and_line(data_near_wall, filepath)

    # 激波区域（dis < 近壁面距离且x在激波区间内）
    data_shock_wave = data[(data["dis"] < config["near_wall_dis"])
                           & (data["X"] < config["shock_wave_x_max"])
                           & (data["X"] > config["shock_wave_x_min"])]
    filepath = os.path.join(prefix, shock_wave + '_error.png')
    plt_hist_and_line(data_shock_wave, filepath)

    # 误差与距离关系分析图
    analyze_path = os.path.join(prefix, train_or_test + '_error_dis_analyze.png')
    plt_error_dis_analyze(config, data, analyze_path)


def plt_hist_and_line(data, filepath):
    """
    绘制四张子图：
    1. 后处理预测值与真实值对比折线图
    2. 后处理相对误差分布直方图
    3. 模型输出预测值与真实值对比折线图
    4. 模型输出相对误差分布直方图
    
    Args:
        data: 包含预测值和真实值的数据框
        filepath: 图片保存路径
    """
    plt.cla()
    plt.figure(figsize=(8, 10), dpi=200)
    data_sample = data.sample(200).sort_values(by="model_label")  # 随机采样200个点
    
    # 子图1：后处理预测值与真实值对比
    plt.subplot(4, 1, 1)
    label_sample = data_sample["post_label"].values
    predict_sample = data_sample["post_predict"].values
    data_num = len(label_sample)

    plt.plot(range(data_num), label_sample, c='g', linestyle='-')    # 真实值（绿色实线）
    plt.plot(range(data_num), predict_sample, c='r', linestyle='--')  # 预测值（红色虚线）
    plt.xlabel('index', fontsize=8)
    plt.ylabel('Mut', fontsize=8)
    plt.title('Postprocess Prediction and Label')

    # 子图2：后处理相对误差分布
    plt.subplot(4, 1, 2)
    label = data["post_label"].values
    predict = data["post_predict"].values
    error = np.abs(label - predict) / label * 100  # 计算相对误差百分比

    gap_data = [5 * x for x in range(0, 20)]  # 误差区间：0-5%, 5-10%, ..., 95-100%
    hist, bin_value = np.histogram(error, gap_data)
    other_data = len(error) - sum(hist)
    hist = np.append(hist, [other_data])
    bin_value = np.append(bin_value, [101])
    hist_prob = [x / len(error) for x in hist]  # 计算各区间样本占比
    plt.bar(bin_value[1:], hist_prob, alpha=1, width=0.8)
    
    for value, prob in zip(bin_value[1:], hist_prob):
        plt.text(value, prob + 0.05, f'{prob:0.2f}', size=8, ha='center', va='bottom')
    plt.xlabel('Relative Error Percent', fontsize=8)
    plt.ylabel('Cell Number Percent', fontsize=8)
    plt.title('Postprocess Relative Error Distribution')

    # 子图3：模型输出预测值与真实值对比
    plt.subplot(4, 1, 3)
    label_sample = data_sample["model_label"].values
    predict_sample = data_sample["model_predict"].values
    data_num = len(label_sample)

    plt.plot(range(data_num), label_sample, c='g', linestyle='-')
    plt.plot(range(data_num), predict_sample, c='r', linestyle='--')
    plt.xlabel('index', fontsize=8)
    plt.ylabel('Mut', fontsize=8)
    plt.title('Model Out Prediction and Label')

    # 子图4：模型输出相对误差分布
    plt.subplot(4, 1, 4)
    model_label = data["model_label"].values
    model_predict = data["model_predict"].values
    error = np.abs(model_label - model_predict) / model_label * 100

    gap_data = [5 * x for x in range(0, 20)]
    hist, bin_value = np.histogram(error, gap_data)
    other_data = len(error) - sum(hist)
    hist = np.append(hist, [other_data])
    bin_value = np.append(bin_value, [101])
    hist_prob = [x / len(error) for x in hist]
    plt.bar(bin_value[1:], hist_prob, alpha=1, width=0.8)
    
    for value, prob in zip(bin_value[1:], hist_prob):
        plt.text(value, prob + 0.05, f'{prob:0.2f}', size=8, ha='center', va='bottom')
    plt.xlabel('Relative Error Percent', fontsize=8)
    plt.ylabel('Cell Number Percent', fontsize=8)
    plt.title('Model Out Relative Error Distribution')

    # 计算并打印评估指标
    condition = re.split('/|_error', filepath)[-2]
    r2_value = r2_score(label, predict)
    print(condition + f":   R2 Score is {r2_value}")

    mse = mean_squared_error(model_label, model_predict)
    print(condition + f":  MSE Score is {mse}")

    plt.tight_layout()
    plt.savefig(filepath)
    plt.clf()


def get_mean_error_by_condition(bins, internal, data):
    """
    按壁面距离区间计算平均误差
    
    Args:
        bins: 区间数量
        internal: 区间长度
        data: 包含error和dis的数据框
        
    Returns:
        tuple: 各区间平均误差列表和区间边界列表
    """
    error_list = []
    for i in range(1, bins + 1):
        # 筛选[internal*i, internal*(i+1))区间内的数据
        data_out = data[(data['dis'] < internal * (i + 1)) & (data['dis'] >= internal * i)]
        error_list.append(data_out['error'].mean())  # 计算区间平均误差
    return error_list, [internal * x for x in range(1, bins + 1)]  # 返回误差列表和区间边界


def get_mean_dis_by_condition(bins, internal, data):
    """
    按误差区间计算平均壁面距离
    
    Args:
        bins: 区间数量
        internal: 区间长度
        data: 包含error和dis的数据框
        
    Returns:
        tuple: 各区间平均壁面距离列表和区间边界列表
    """
    error_list = []
    for i in range(1, bins + 1):
        # 筛选[internal*i, internal*(i+1))区间内的数据
        data_out = data[(data['error'] < internal * (i + 1)) & (data['error'] >= internal * i)]
        error_list.append(data_out['dis'].mean())  # 计算区间平均壁面距离
    return error_list, [internal * x for x in range(1, bins + 1)]  # 返回距离列表和区间边界


def format_sci(data_list):
    """
    将数据列表格式化为科学计数法字符串
    
    Args:
        data_list: 数值列表
        
    Returns:
        list: 科学计数法字符串列表
    """
    out = []
    for data in data_list:
        out.append(f"{data:0.1E}")  # 保留1位小数的科学计数法
    return out